#pragma once

#include <cstdint>
#include <memory>
#include <string>
#include <stdexcept>
#include <vector>

#include "infer_config.h"

namespace infer {
// TODO 待删除, 使用 txnncpp struct Dims
using dim_t               = std::int64_t;

// TODO 待修改, 使用 txnncpp 类型
using ITMExecutionContext = std::string;
using ITMEngine           = std::string;
using ITMRuntime          = std::string;

/**
 * @brief TxNN后端管理类, 对后端具体推理操作的封装。
 */
class TxnnManager {
public:
  TxnnManager(const std::string& engine_path, const InferConfig& cfg);
  TxnnManager();
  ~TxnnManager();

  TxnnManager(const TxnnManager&)             = delete; // 禁用拷贝构造函数
  TxnnManager& operator=(const TxnnManager&)  = delete; // 禁用拷贝赋值运算符
  TxnnManager(TxnnManager&&)                  = delete; // 禁用移动构造函数
  TxnnManager& operator=(TxnnManager&&)       = delete; // 禁用移动赋值运算符

  /**
   * @brief 初始化方法，用于加载 txnn 引擎(编译产物)。
   * @param blob 包含引擎数据的指针。
   * @param size 引擎数据的大小。
   */
  void init(const void* blob, std::size_t size);

  /**
   * @brief 设置张量地址。
   * @param tensor_name 张量名字。
   * @param data 张量数据指针。
   * @return bool 设置是否成功。
   */
  bool setTensorAddress(const char* tensor_name, void* data);

  /**
   * @brief 设置输入张量形状。
   * @param tensor_name 张量名字。
   * @param dims 张量维度。
   * @return bool 设置是否成功。
   */
  bool setInputShape(const char* tensor_name, const std::vector<dim_t>& dims);

  /**
   * @brief 指定流上执行推理。
   * @param stream TxDA Stream, 目前不可用。// TODO TxDA 框架添加
   * @return bool 设置是否成功。
   */
  bool enqueue();

  /**
   * @brief 获取张量形状。
   * @param tensor_name 张量名字。
   * @param dims 张量维度。
   * @return bool 设置是否成功。
   */
  std::vector<dim_t> getTensorShape(const char* tensor_name) const noexcept;

  /**
   * @brief 获取张量的输入输出模式。
   * @param tensor_name 张量名字。
   * @param dims 张量维度。
   * @return std::int32_t 张量的输入输出模式: 0-不是输入/输出, 1-输入, 2-输出。
   */
  // TODO 待删除, 使用 txnncpp enum class TensorIOMode : std::int32_t
  std::int32_t getTensorIOMode(const char* tensor_name) const noexcept;

  /**
   * @brief 获取输入输出张量的数量。
   * @return std::int32_t 输出张量的数量。
   */
  // TODO 待删除, 使用 txnncpp enum class TensorIOMode : std::int32_t
  std::int32_t getNbIOTensors() const noexcept;
  
  /**
   * @brief 获取指定索引的输入输出张量的名字。
   * @param index 张量索引。
   * @return std::int32_t 输出张量的数量。
   */
  const char* getIOTensorName(std::int32_t index) const noexcept;

private:
  std::string engine_path_; // TxnnCPP 引擎路径
  InferConfig cfg_;         // TxnnCPP 配置

  std::unique_ptr<ITMExecutionContext>  context_; // TxnnCPP 执行上下文
  std::shared_ptr<ITMEngine>            engine_;  // TxnnCPP TxDA引擎, 是否可共享?
  std::unique_ptr<ITMRuntime>           runtime_; // TxnnCPP 运行时
};

} // namespace infer